import sqlite3
import uuid
import re
import pytz

from datetime import date, datetime

from typing import Optional
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import tool
from langgraph_customer.utils.Retriever import VectorStoreRetriever
from langchain_core.documents import Document

# 创建数据库连接
db = "../resource/travel2.sqlite"

@tool
def lookup_policy(query: str, docs: list) -> str:
    """
    通过向量检索器，查询公司政策知识库
    """
    retriever = VectorStoreRetriever()
    retriever.put_documents(docs)
    return_docs = retriever.get_similar_docs(query)
    # 将结果拼接
    return "\n".join([doc.page_content for doc in return_docs])

@tool
def fetch_user_flight_information(config: RunnableConfig) -> list[dict]:
    """
    获取用户的所有机票信息以及相应的航班信息和座位分配

    返回：
        一个字典列表，其中每个字典包含机票详情
        相关航班详情以及用户每张机票的座位分配
    """
    configuration = config.get("configurable", {})
    passenger_id = configuration.get("passenger_id", None)
    if not passenger_id:
        raise ValueError("未配置passenger_id")

    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    query = """
    SELECT 
        t.ticket_no, t.book_ref,
        f.flight_id, f.flight_no, f.departure_airport, f.arrival_airport, f.scheduled_departure, f.scheduled_arrival,
        bp.seat_no, tf.fare_conditions
    FROM 
        tickets t
        JOIN ticket_flights tf ON t.ticket_no = tf.ticket_no
        JOIN flights f ON tf.flight_id = f.flight_id
        JOIN boarding_passes bp ON bp.ticket_no = t.ticket_no AND bp.flight_id = f.flight_id
    WHERE 
        t.passenger_id = ?
    """
    cursor.execute(query, (passenger_id,))
    rows = cursor.fetchall()
    column_names = [column[0] for column in cursor.description]
    results = [dict(zip(column_names, row)) for row in rows]

    cursor.close()
    conn.close()

    return results

@tool
def search_flights(
    departure_airport: Optional[str] = None,
    arrival_airport: Optional[str] = None,
    start_time: Optional[date | datetime] = None,
    end_time: Optional[date | datetime] = None,
    limit: int = 20,
) -> list[dict]:
    """
    根据出发机场、到达机场和出发时间范围搜索航班
    """
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    query = "SELECT * FROM flights WHERE 1 = 1"
    params = []
    # 添加条件
    if departure_airport:
        query += " AND departure_airport = ?"
        params.append(departure_airport)

    if arrival_airport:
        query += " AND arrival_airport = ?"
        params.append(arrival_airport)

    if start_time:
        query += " AND scheduled_departure >= ?"
        params.append(start_time)

    if end_time:
        query += " AND scheduled_departure <= ?"
        params.append(end_time)
    query += " LIMIT ?"
    params.append(limit)
    cursor.execute(query, params)
    rows = cursor.fetchall()
    column_names = [column[0] for column in cursor.description]
    results = [dict(zip(column_names, row)) for row in rows]

    cursor.close()
    conn.close()

    return results


@tool
def update_ticket_to_new_flight(
    ticket_no: str, new_flight_id: int, *, config: RunnableConfig
) -> str:
    """
    将用户的机票更新为新的有效航班
    """
    configuration = config.get("configurable", {})
    passenger_id = configuration.get("passenger_id", None)
    if not passenger_id:
        raise ValueError("未配置passenger_id")

    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    cursor.execute(
        "SELECT departure_airport, arrival_airport, scheduled_departure FROM flights WHERE flight_id = ?",
        (new_flight_id,),
    )
    new_flight = cursor.fetchone()
    if not new_flight:
        cursor.close()
        conn.close()
        return "提供的new_flight_id无效。"
    column_names = [column[0] for column in cursor.description]
    new_flight_dict = dict(zip(column_names, new_flight))
    # 更改时区
    timezone = pytz.timezone("Etc/GMT-3")
    current_time = datetime.now(tz=timezone)
    departure_time = datetime.strptime(
        new_flight_dict["scheduled_departure"], "%Y-%m-%d %H:%M:%S.%f%z"
    )
    time_until = (departure_time - current_time).total_seconds()
    if time_until < (3 * 3600):
        return f"不允许改期至距当前时间少于 3 小时的航班。所选航班{departure_time}."

    cursor.execute(
        "SELECT flight_id FROM ticket_flights WHERE ticket_no = ?", (ticket_no,)
    )
    current_flight = cursor.fetchone()
    if not current_flight:
        cursor.close()
        conn.close()
        return "未找到与给定票号对应的现有票。"

    # 检查登录用户是否确实拥有此票
    cursor.execute(
        "SELECT * FROM tickets WHERE ticket_no = ? AND passenger_id = ?",
        (ticket_no, passenger_id),
    )
    current_ticket = cursor.fetchone()
    if not current_ticket:
        cursor.close()
        conn.close()
        return f"当前已登录的passenger_id为 {passenger_id}，不是ticket_no为 {ticket_no} 的所有者"

    # 在实际应用中，您可能会在此处添加额外的检查来强制执行业务逻辑，例如“新的出发机场是否与当前机票匹配”等等。
    # 虽然最好尝试*主动*地向 LLM 提供“类型提示”策略，这不可避免地会出错，因此您还需要确保您的API强制执行有效的行为
    cursor.execute(
        "UPDATE ticket_flights SET flight_id = ? WHERE ticket_no = ?",
        (new_flight_id, ticket_no),
    )
    conn.commit()

    cursor.close()
    conn.close()
    return "机票已成功更新为新航班。"


@tool
def cancel_ticket(ticket_no: str, *, config: RunnableConfig) -> str:
    """
    取消用户的票并将其从数据库中删除。
    """
    configuration = config.get("configurable", {})
    passenger_id = configuration.get("passenger_id", None)
    if not passenger_id:
        raise ValueError("未配置passenger_id")
    conn = sqlite3.connect(db)
    cursor = conn.cursor()

    cursor.execute(
        "SELECT flight_id FROM ticket_flights WHERE ticket_no = ?", (ticket_no,)
    )
    existing_ticket = cursor.fetchone()
    if not existing_ticket:
        cursor.close()
        conn.close()
        return "未找到与给定票号对应的现有票。"

    # Check the signed-in user actually has this ticket
    cursor.execute(
        "SELECT ticket_no FROM tickets WHERE ticket_no = ? AND passenger_id = ?",
        (ticket_no, passenger_id),
    )
    current_ticket = cursor.fetchone()
    if not current_ticket:
        cursor.close()
        conn.close()
        return f"当前已登录的passenger_id 为 {passenger_id}，不是ticket_no为 {ticket_no} 的所有者"

    cursor.execute("DELETE FROM ticket_flights WHERE ticket_no = ?", (ticket_no,))
    conn.commit()

    cursor.close()
    conn.close()
    return "票已成功取消。"



if __name__ == '__main__':
    # 加载航班订票相关知识数据
    with open("../resource/swiss_faq.md", "r", encoding="utf-8") as f:
        text = f.read()

    docs = [txt for txt in re.split(r"(?=##)", text)]
    docs.pop(0)
    new_docs = [Document(page_content=doc, id=str(uuid.uuid4())) for doc in docs]
    data = lookup_policy("如何申请退票？", new_docs)

    #     thread_id = str(uuid.uuid4())
    #     config = {
    #         "configurable": {
    #             # 航班工具会使用 passenger_id 来获取用户的航班信息
    #             "passenger_id": "3442 587242",
    #             # 检查点通过thread_id访问
    #             "thread_id": thread_id,
    #         }
    #     }
    #
    #     data = fetch_user_flight_information(config, "../resource/travel2.sqlite")
    #     print(data)
